# 弹簧阻尼系统的LQR控制
# xd[k+1] = ADXd[k],AD目标转移矩阵，为单位阵时，追踪的目标为常数
# 是调用<F1_LQR_Gain>函数

import F_LQR_Gain 
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
import control as ct
import time
import scipy.linalg as la
from scipy.signal import StateSpace

# 1.定义系统
m = 1
k = 1
b = 0.5
# x[k+1] = Ax[k] + Bu[k]
A = np.array([[0,1],[-k/m,-b/m]])
B = np.array([[0],[1/m]])
C = np.array([[1,0]])
D = 0
# 创建状态空间模型
# sys = StateSpace(A, B, C, D)
Ts = 0.1
sys_continuous = ct.StateSpace(A, B, C, D)
sys_discrete = ct.sample_system(sys_continuous, Ts, method='zoh')
# 获取行数
n = len(A)
# 获取列数
p =  B.shape[1]
# 离散化后的状态矩阵和输入矩阵
A = sys_discrete.A
B = sys_discrete.B
# 目标
xd = np.array([[1],[0]]).reshape(1,-1)
AD = np.eye(n)
# 2.初始化系统
x0 = np.array([[0],[0]]).reshape(1,-1)
x = x0
xa = np.hstack((x,xd))
# 输入
u0 = np.array([[0]]).reshape(1,-1)
u = u0

# 定义系统运行步数
k_steps = 100
x_history = np.zeros((n,k_steps+1))
x_history[:,0] = x
u_history = np.zeros((p,k_steps))
u_history[:,0] = u

# 设置权重
Q = np.array([[1,0],[0,1]])
S = np.array([[1,0],[0,1]])
R=0.1
N = k_steps
P_k = S
# 组合新的矩阵
Ca = np.hstack((np.eye(n),-np.eye(n)))#横向堆叠
Aa = np.vstack((np.hstack((A,np.zeros((n,n)))),np.hstack((np.zeros((n,n)),AD))))
Ba = np.vstack((B,np.zeros((n,p))))
Sa = Ca.T@S@Ca
Qa = Ca.T@Q@Ca
# 3.计算反馈增益
F = F_LQR_Gain.LQR_Gain(Aa,Ba,Qa,R,Sa)

for k in range(1,k_steps+1):
    u = -F@xa.T
    x=A@x_history[:,k-1].reshape(-1,1)+B@u
    xa = np.vstack((x,xd.T)).T
    x_history[:,k]=x.T
    u_history[:,k-1]=u
 
"""
# # 写入数据到表格中
from openpyxl import Workbook
# 创建 Excel 工作簿
wb = Workbook()
ws = wb.active
# 将矩阵数据写入工作表
for row in x_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('x_history.xlsx')
wb = Workbook()
ws = wb.active

# 将矩阵数据写入工作表
for row in u_history:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('u_history.xlsx')

# 创建 Excel 工作簿
wb = Workbook()
ws = wb.active
# 将矩阵数据写入工作表
for row in F_N:
    ws.append(row.tolist())
# 保存 Excel 文件
wb.save('F_N.xlsx')

""" 


# 4.制图
plt.figure()
plt.subplot(2,1,1)
plt.title('Input and State Plot')
plt.xlabel('steps')
plt.ylabel('State')
for i in range(0,n): 
    plt.plot(x_history[i,:],label=f'State {i+1}')
    
plt.legend()
plt.subplot(2,1,2)
plt.ylabel('Input')
for i in range(0,p):
    plt.plot(u_history[i,:],label=f'Input {i+1}')
    
plt.legend()
plt.grid()
plt.show()